import asyncio
import os
import dotenv

from openai import AsyncOpenAI
from openai.types.responses import ResponseTextDeltaEvent

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    function_tool,
    set_tracing_disabled,
)

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)

CUSTOM_MODEL_PROVIDER = CustomModelProvider()

async def main():
    agent = Agent(
        name="王富贵",
        instructions="""
        你是一个agent，你是一个有帮助的助手。       
        """,
    )
    result = Runner.run_streamed(agent, input="请给我讲5个成人笑话。", run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),)

    async for event in result.stream_events():
        if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
            print(event.data.delta, end="", flush=True)

if __name__ == "__main__":
    asyncio.run(main())